import os
import pandas as pd
from settings import DATA_PATH, PROCESSED_DATA_PATH
import numpy as np


def load_book_info():
    """原始数据集"""
    path = os.path.join(DATA_PATH, 'books/book.csv')
    return pd.read_csv(path, sep='\t')


def load_book_clean():
    """经过预处理后的数据集"""
    path = os.path.join(PROCESSED_DATA_PATH, 'book_clean.pkl')
    df = pd.read_pickle(path)
    return df


def load_book_augmented():
    """经过文本增强后的数据集"""
    path = os.path.join(PROCESSED_DATA_PATH, 'book_augmented.pkl')
    df = pd.read_pickle(path)
    return df


def load_label_dict():
    """类别-类别索引的字典"""
    path = os.path.join(PROCESSED_DATA_PATH, 'label_dict.pkl')
    return pd.read_pickle(path)


def load_book_undersample_1k():
    """下采样的一个小型数据集，每个类别都有1000条样本"""
    path = os.path.join(PROCESSED_DATA_PATH, 'undersample_1k.pkl')
    return pd.read_pickle(path)


class FeatureContainer:
    """
    用于特征保存、生成pd.DataFrame与加载
    """
    @staticmethod
    def save(book_ids, features, name):
        """
        存储特征
        :param book_ids: 书籍id，np.ndarry，必须为2维，N x 1
        :param features:  特征，np.ndarry
        :param name: 文件名，无需加.pkl
        """
        columns = ['book_id'] + [name + '_{}'.format(i) for i in range(features.shape[1])]
        data = pd.DataFrame(np.concatenate([book_ids, features], axis=1), columns=columns)
        data['book_id'] = data['book_id'].astype(int)
        print('已保存特征{}'.format(name))
        data.to_pickle(os.path.join(PROCESSED_DATA_PATH, name + '.pkl'))
        
    @staticmethod
    def make_df(book_ids, features, name):
        """
        生成pd.DataFrame
        :param book_ids: 书籍id，np.ndarry，必须为2维，N x 1
        :param features:  特征，np.ndarry
        :param name: 特征列名前缀
        """
        columns = ['book_id'] + [name + '_{}'.format(i) for i in range(features.shape[1])]
        data = pd.DataFrame(np.concatenate([book_ids, features], axis=1), columns=columns)
        data['book_id'] = data['book_id'].astype(int)
        return data
    
    @staticmethod
    def load(names, labels):
        """
        加载特征
        :param name: 特征文件名
        :param labels: 文本的标签，pd.DataFrame，列名：'book_id','label'。最终抽取的样本，以labels中的book_id为准
        :return: merge后的完整数据集
        """
        book_ids = []
        data_list = []
        if len(names) == 0 or isinstance(names, str):
            features = pd.read_pickle(os.path.join(PROCESSED_DATA_PATH, names + '.pkl'))
            return pd.merge(features, labels, how='inner', on='book_id')
        else:
            for name in names:
                data = pd.read_pickle(os.path.join(PROCESSED_DATA_PATH, name + '.pkl'))
                book_ids.extend(data['book_id'].tolist())
                data_list.append(data)
            book_ids = pd.DataFrame(np.unique(book_ids).reshape(-1,1), columns=['book_id'])
        
            features = None
            for i, data in enumerate(data_list):
                if i == 0:
                    features = pd.merge(book_ids, data, how='left', on='book_id')
                else:
                    features = pd.merge(features, data, how='left', on='book_id')
            features = pd.merge(features, labels, how='inner', on='book_id')
            return features